#include <cuda_runtime.h>
#include <stdio.h>

#include <iostream>

#include "common/config.h"
#include "gpu/kernel.h"

__device__ __host__ float normalize_angle(float rad) {
    const float PI = 3.14159265358979323846f;
    const float TWO_PI = 2.0f * PI;
    const float INV_TWO_PI = 1.0f / TWO_PI;

    // 归一化到 [-1, 1]（对应 [-π, π]）
    rad = rad * INV_TWO_PI;
    rad = rad - floorf(rad + 0.5f);  // 等价于 fmodf(rad + PI, TWO_PI) - PI
    return rad * TWO_PI;
}

__global__ void spatical_sampling_kernel(char* data, float min_ref_s, float max_ref_s, int refline_offset, int path_info_offset) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    int by = blockIdx.y;
    float s = idx * kPointInfoResolution;  // 需要计算s处的point_info

    PathInfo* path_info = reinterpret_cast<PathInfo*>(data + path_info_offset) + by;

    if (s > path_info->max_s || s < path_info->min_s || idx >= path_info->point_info_nums) {
        return;
    }
    CubicPolynomialCoeffs* path_coeffs = reinterpret_cast<CubicPolynomialCoeffs*>(data + path_info->cubic_polynomial_coeff_offset);
    int point_coeff_nums = path_info->cubic_polynomial_coeff_nums;

    int index = 0;
    while (path_coeffs->end_s < s && index < point_coeff_nums) {
        ++path_coeffs;
        ++index;
    }

    float s_squre = s * s;
    float l = path_coeffs->a0 + path_coeffs->a1 * s + path_coeffs->a2 * s_squre + path_coeffs->a3 * s_squre * s;

    PointInfo* point_info = reinterpret_cast<PointInfo*>(data + path_info->point_info_offset) + idx;
    point_info->path_index = path_info->path_index;
    point_info->point_index = idx;
    point_info->s = s;
    point_info->l = l;

    RefPoint* ref_point = reinterpret_cast<RefPoint*>(data + refline_offset) + idx;
    point_info->theta_ref = ref_point->theta;
    float angle_dir = l > 0 ? ref_point->theta + M_PI / 2 : ref_point->theta - M_PI / 2;
    angle_dir = normalize_angle(angle_dir);
    float l_abs = fabsf(l);

    point_info->x = ref_point->x + l_abs * cosf(angle_dir);
    point_info->y = ref_point->y + l_abs * sinf(angle_dir);

    // if (idx == 100) {  // debug
    //     printf("blockIdx:(%d, %d), threadIdx:(%d, %d), path_index: %d, point_info_nums: %d, point_info_offset: %d, point_info.s: %f, point_info.l: %f\n",
    //     blockIdx.x, blockIdx.y, threadIdx.x, threadIdx.y,
    //            path_info->path_index, path_info->point_info_nums, path_info->point_info_offset, point_info->s, point_info->l);
    // }
}

extern "C" void spatical_sampling_launcher(cudaStream_t* stream, const DataInfo* data_info, char* data) {
    float max_length = data_info->max_length;
    int path_nums = data_info->path_nums;
    int max_point_info_nums = int((max_length + 1e-6) / kPointInfoResolution);
    int block_x = 128;
    int block_y = 1;
    dim3 block_size(block_x, block_y);

    int grid_x = static_cast<int>((max_point_info_nums + block_x - 1) / block_x);
    int grid_y = path_nums;
    dim3 grid_size(grid_x, grid_y);

    spatical_sampling_kernel<<<grid_size, block_size, 0, *stream>>>(data, data_info->min_ref_s, data_info->max_ref_s, data_info->refline_offset,
                                                                    data_info->path_info_offset);
    return;
}